#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
任务分发器
统一处理来自不同源的任务（HTTP API、远程调度系统），并分发给现有的工作流管理器
"""

import logging
import uuid
from typing import Dict, Any, Optional
from datetime import datetime
from enum import Enum

try:
    # 尝试相对导入（当作为模块使用时）
    from .scheduler_connector import RemoteTask
    from .workflow_manager import WorkflowManager
    from .export_controller import ExportResolution, ExportFramerate
    from .com_thread_manager import get_com_thread_manager, TaskPriority
    from .config_manager import get_config
except ImportError:
    # 直接导入（当直接运行时）
    from scheduler_connector import RemoteTask
    from workflow_manager import WorkflowManager
    from export_controller import ExportResolution, ExportFramerate
    from com_thread_manager import get_com_thread_manager, TaskPriority
    from config_manager import get_config


class TaskSource(Enum):
    """任务来源枚举"""
    HTTP_API = "http_api"
    REMOTE_SCHEDULER = "remote_scheduler"


class TaskDispatcher:
    """任务分发器"""

    def __init__(self):
        """初始化任务分发器"""
        self.logger = logging.getLogger(__name__)
        self.config = get_config()

        # 初始化组件
        self.workflow_manager = WorkflowManager()
        self.com_manager = get_com_thread_manager()

        # 任务映射：用于跟踪远程任务ID到内部任务ID的映射
        self.remote_task_mapping: Dict[str, str] = {}

        # 回调函数
        self.progress_callback: Optional[callable] = None
        self.completion_callback: Optional[callable] = None

        self.logger.info("任务分发器初始化完成")

    def set_progress_callback(self, callback: callable):
        """设置进度回调函数"""
        self.progress_callback = callback

    def set_completion_callback(self, callback: callable):
        """设置完成回调函数"""
        self.completion_callback = callback

    def dispatch_remote_task(self, remote_task: RemoteTask) -> Dict[str, Any]:
        """分发远程调度任务"""
        try:
            self.logger.info(f"开始分发远程任务: {remote_task.id}, 类型: {remote_task.taskType}")

            # 解析远程任务数据
            task_params = self._parse_remote_task(remote_task)
            if not task_params:
                error_msg = f"无法解析远程任务数据: {remote_task.id}"
                self.logger.error(error_msg)
                return {
                    "success": False,
                    "message": error_msg,
                    "remote_task_id": remote_task.id
                }

            # 生成内部任务ID
            internal_task_id = f"remote_{remote_task.id}_{uuid.uuid4().hex[:8]}"

            # 建立映射关系
            self.remote_task_mapping[remote_task.id] = internal_task_id

            # 创建任务包装器
            def wrapped_workflow_task(**kwargs):
                """包装的工作流任务，用于处理进度和完成回调"""
                return self._execute_workflow_with_callbacks(
                    remote_task.id,
                    internal_task_id,
                    task_params
                )

            # 提交到COM线程管理器
            self.com_manager.submit_task(
                internal_task_id,
                wrapped_workflow_task,
                args=(),
                kwargs=task_params,
                priority=TaskPriority.HIGH  # 远程任务使用高优先级
            )

            self.logger.info(f"远程任务已提交: {remote_task.id} -> {internal_task_id}")

            return {
                "success": True,
                "message": "任务已提交处理",
                "remote_task_id": remote_task.id,
                "internal_task_id": internal_task_id,
                "timestamp": datetime.now().isoformat()
            }

        except Exception as e:
            error_msg = f"分发远程任务异常: {str(e)}"
            self.logger.error(error_msg)
            return {
                "success": False,
                "message": error_msg,
                "remote_task_id": remote_task.id,
                "timestamp": datetime.now().isoformat()
            }

    def dispatch_http_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """分发HTTP API任务"""
        try:
            self.logger.info(f"开始分发HTTP任务: {task_data.get('url', 'unknown')}")

            # 验证必要参数
            if 'url' not in task_data:
                return {
                    "success": False,
                    "message": "缺少必需的参数 'url'"
                }

            # 解析HTTP任务参数
            task_params = self._parse_http_task(task_data)

            # 生成任务ID
            task_id = task_data.get('task_id', f"http_{uuid.uuid4().hex[:8]}")

            # 创建任务包装器
            def wrapped_workflow_task(**kwargs):
                """包装的工作流任务"""
                return self.workflow_manager.execute_workflow(**task_params)

            # 提交到COM线程管理器
            self.com_manager.submit_task(
                task_id,
                wrapped_workflow_task,
                args=(),
                kwargs=task_params,
                priority=TaskPriority.NORMAL  # HTTP任务使用普通优先级
            )

            self.logger.info(f"HTTP任务已提交: {task_id}")

            return {
                "success": True,
                "message": "任务已提交处理",
                "task_id": task_id,
                "timestamp": datetime.now().isoformat()
            }

        except Exception as e:
            error_msg = f"分发HTTP任务异常: {str(e)}"
            self.logger.error(error_msg)
            return {
                "success": False,
                "message": error_msg,
                "timestamp": datetime.now().isoformat()
            }

    def _parse_remote_task(self, remote_task: RemoteTask) -> Optional[Dict[str, Any]]:
        """解析远程任务数据"""
        try:
            # 获取任务参数，优先使用taskParams
            task_params_data = remote_task.taskParams

            # 如果taskParams为空，尝试向后兼容的cozeDraftData
            if not task_params_data:
                task_params_data = remote_task.cozeDraftData

            if not task_params_data:
                self.logger.error(f"远程任务缺少参数数据: {remote_task.id}")
                return None

            # 根据任务类型解析不同的参数格式
            if remote_task.taskType == "VIDEO_RENDER":
                return self._parse_video_render_params(remote_task, task_params_data)
            elif remote_task.taskType == "IMAGE_PROCESSING":
                return self._parse_image_processing_params(remote_task, task_params_data)
            else:
                self.logger.warning(f"未知任务类型: {remote_task.taskType}")
                # 回退到通用解析
                return self._parse_generic_params(remote_task, task_params_data)

        except Exception as e:
            self.logger.error(f"解析远程任务数据异常: {e}")
            return None

    def _parse_video_render_params(self, remote_task: RemoteTask, task_params_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """解析VIDEO_RENDER任务参数"""
        try:
            # 提取必要的URL
            url = task_params_data.get('draft_url') or task_params_data.get('url') or task_params_data.get('draftUrl')
            if not url:
                self.logger.error(f"VIDEO_RENDER任务缺少draft_url: {remote_task.id}")
                return None

            # 使用配置的默认输出路径
            output_path = self.config.get_output_path()

            # 提取草稿名称
            draft_name = (task_params_data.get('draft_name') or
                         task_params_data.get('draftName') or
                         task_params_data.get('name') or
                         self._extract_draft_name_from_url(url))

            # 解析分辨率和帧率
            resolution = self._parse_resolution(task_params_data.get('resolution'))
            framerate = self._parse_framerate(task_params_data.get('framerate'))

            task_params = {
                'url': url,
                'output_path': output_path,
                'draft_name': draft_name,
                'resolution': resolution,
                'framerate': framerate,
                'task_id': remote_task.id
            }

            self.logger.debug(f"解析VIDEO_RENDER任务参数: {task_params}")
            return task_params

        except Exception as e:
            self.logger.error(f"解析VIDEO_RENDER任务参数异常: {e}")
            return None

    def _parse_image_processing_params(self, remote_task: RemoteTask, task_params_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """解析IMAGE_PROCESSING任务参数"""
        try:
            # 预留IMAGE_PROCESSING任务的参数解析逻辑
            self.logger.info(f"IMAGE_PROCESSING任务参数解析待实现: {remote_task.id}")

            # 目前返回基本参数
            task_params = {
                'task_id': remote_task.id,
                'raw_params': task_params_data
            }

            return task_params

        except Exception as e:
            self.logger.error(f"解析IMAGE_PROCESSING任务参数异常: {e}")
            return None

    def _parse_generic_params(self, remote_task: RemoteTask, task_params_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """通用参数解析（向后兼容）"""
        try:
            # 提取必要的URL
            url = task_params_data.get('url') or task_params_data.get('draft_url') or task_params_data.get('draftUrl')
            if not url:
                self.logger.error(f"远程任务缺少URL: {remote_task.id}")
                return None

            # 使用配置的默认输出路径
            output_path = self.config.get_output_path()

            # 提取草稿名称
            draft_name = (task_params_data.get('draft_name') or
                         task_params_data.get('draftName') or
                         task_params_data.get('name') or
                         self._extract_draft_name_from_url(url))

            # 解析分辨率和帧率
            resolution = self._parse_resolution(task_params_data.get('resolution'))
            framerate = self._parse_framerate(task_params_data.get('framerate'))

            task_params = {
                'url': url,
                'output_path': output_path,
                'draft_name': draft_name,
                'resolution': resolution,
                'framerate': framerate,
                'task_id': remote_task.id
            }

            self.logger.debug(f"解析通用任务参数: {task_params}")
            return task_params

        except Exception as e:
            self.logger.error(f"解析通用任务参数异常: {e}")
            return None

    def _parse_http_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """解析HTTP任务数据"""
        try:
            url = task_data['url'].strip()

            # 使用配置的默认输出路径
            output_path = task_data.get('output_path') or self.config.get_output_path()

            # 提取草稿名称
            draft_name = (task_data.get('draft_name') or
                         self._extract_draft_name_from_url(url))

            # 解析分辨率和帧率
            resolution = self._parse_resolution(task_data.get('resolution'))
            framerate = self._parse_framerate(task_data.get('framerate'))

            task_params = {
                'url': url,
                'output_path': output_path,
                'draft_name': draft_name,
                'resolution': resolution,
                'framerate': framerate,
                'task_id': task_data.get('task_id')
            }

            self.logger.debug(f"解析HTTP任务参数: {task_params}")
            return task_params

        except Exception as e:
            self.logger.error(f"解析HTTP任务数据异常: {e}")
            raise

    def _execute_workflow_with_callbacks(self, remote_task_id: str, internal_task_id: str, task_params: Dict[str, Any]) -> Dict[str, Any]:
        """执行工作流并处理回调"""
        try:
            self.logger.info(f"开始执行工作流 - 远程任务: {remote_task_id}")

            # 报告任务开始处理
            if self.progress_callback:
                try:
                    self.progress_callback(remote_task_id, 0)
                except Exception as e:
                    self.logger.warning(f"进度回调异常: {e}")

            # 执行工作流
            result = self.workflow_manager.execute_workflow(**task_params)

            # 处理结果
            if result.get("status") == "completed":
                self.logger.info(f"工作流执行成功 - 远程任务: {remote_task_id}")

                # 报告任务完成
                if self.completion_callback:
                    try:
                        # 准备完成数据
                        completion_data = {
                            "result": result.get("result", {}),
                            "final_path": result.get("result", {}).get("final_path"),
                            "timestamp": datetime.now().isoformat()
                        }

                        # 如果有OSS上传结果，添加OSS URL
                        oss_upload = result.get("oss_upload", {})
                        oss_video_url = oss_upload.get("oss_url") if oss_upload.get("success") else None

                        self.completion_callback(remote_task_id, completion_data, oss_video_url)
                    except Exception as e:
                        self.logger.error(f"完成回调异常: {e}")
                        # 即使回调失败，也不影响工作流结果
            else:
                self.logger.error(f"工作流执行失败 - 远程任务: {remote_task_id}")

                # 报告任务失败
                if self.completion_callback:
                    try:
                        error_message = result.get("message", "工作流执行失败")
                        # 使用错误消息作为特殊标记来区分失败情况
                        self.completion_callback(remote_task_id, None, None, error_message)
                    except Exception as e:
                        self.logger.error(f"失败回调异常: {e}")

            # 清理映射
            if remote_task_id in self.remote_task_mapping:
                del self.remote_task_mapping[remote_task_id]

            return result

        except Exception as e:
            error_msg = f"工作流执行异常: {str(e)}"
            self.logger.error(error_msg)

            # 报告任务失败
            if self.completion_callback:
                try:
                    self.completion_callback(remote_task_id, None, None, error_msg)
                except Exception as callback_e:
                    self.logger.error(f"异常回调失败: {callback_e}")

            # 清理映射
            if remote_task_id in self.remote_task_mapping:
                del self.remote_task_mapping[remote_task_id]

            return {
                "status": "failed",
                "error": error_msg,
                "task_id": internal_task_id,
                "timestamp": datetime.now().isoformat()
            }

    def _extract_draft_name_from_url(self, url: str) -> Optional[str]:
        """从URL中提取草稿名称"""
        try:
            import re
            # 查找 .json 前的UUID模式
            pattern = r'/([a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12})\.json'
            match = re.search(pattern, url)

            if match:
                return match.group(1)

            # 备选方案：查找draft/后面的文件名
            alt_pattern = r'/draft/([^/]+?)(?:\.json)?(?:\?|$)'
            alt_match = re.search(alt_pattern, url)

            if alt_match:
                draft_id = alt_match.group(1)
                if draft_id.endswith('.json'):
                    draft_id = draft_id[:-5]
                return draft_id

            return None

        except Exception as e:
            self.logger.warning(f"提取草稿名异常: {e}")
            return None

    def _parse_resolution(self, resolution_str: Optional[str]) -> Optional[ExportResolution]:
        """解析分辨率字符串"""
        if not resolution_str:
            return None

        try:
            return ExportResolution(resolution_str)
        except ValueError:
            self.logger.warning(f"无效的分辨率: {resolution_str}")
            return None

    def _parse_framerate(self, framerate_str: Optional[str]) -> Optional[ExportFramerate]:
        """解析帧率字符串"""
        if not framerate_str:
            return None

        try:
            return ExportFramerate(framerate_str)
        except ValueError:
            self.logger.warning(f"无效的帧率: {framerate_str}")
            return None

    def get_task_status(self, task_id: str, is_remote: bool = False) -> Optional[Dict[str, Any]]:
        """获取任务状态"""
        try:
            if is_remote:
                # 远程任务：通过映射找到内部任务ID
                internal_task_id = self.remote_task_mapping.get(task_id)
                if not internal_task_id:
                    return None

                # 首先从COM管理器获取状态
                com_status = self.com_manager.get_task_status(internal_task_id)
                if "error" not in com_status:
                    return com_status

                # 回退到工作流管理器
                return self.workflow_manager.get_task_status(internal_task_id)
            else:
                # HTTP任务：直接使用任务ID
                # 首先从COM管理器获取状态
                com_status = self.com_manager.get_task_status(task_id)
                if "error" not in com_status:
                    return com_status

                # 回退到工作流管理器
                return self.workflow_manager.get_task_status(task_id)

        except Exception as e:
            self.logger.error(f"获取任务状态异常: {e}")
            return None

    def cancel_task(self, task_id: str, is_remote: bool = False) -> bool:
        """取消任务"""
        try:
            if is_remote:
                # 远程任务：通过映射找到内部任务ID
                internal_task_id = self.remote_task_mapping.get(task_id)
                if not internal_task_id:
                    return False

                # 从COM管理器和工作流管理器中取消
                com_success = self.com_manager.cancel_task(internal_task_id)
                workflow_success = self.workflow_manager.cancel_task(internal_task_id)

                # 清理映射
                if task_id in self.remote_task_mapping:
                    del self.remote_task_mapping[task_id]

                return com_success or workflow_success
            else:
                # HTTP任务：直接使用任务ID
                com_success = self.com_manager.cancel_task(task_id)
                workflow_success = self.workflow_manager.cancel_task(task_id)
                return com_success or workflow_success

        except Exception as e:
            self.logger.error(f"取消任务异常: {e}")
            return False

    def list_tasks(self) -> Dict[str, Any]:
        """列出所有任务"""
        try:
            # 获取COM管理器的任务
            com_tasks = self.com_manager.list_tasks(include_completed=True)
            self.logger.info(f"📊 COM管理器任务数量: {len(com_tasks)}")
            if com_tasks:
                self.logger.info(f"📊 COM任务ID列表: {list(com_tasks.keys())}")
                for task_id, task_info in list(com_tasks.items())[:3]:  # 只显示前3个
                    self.logger.info(f"   - {task_id}: status={task_info.get('status')}")

            # 获取工作流管理器的任务
            workflow_data = self.workflow_manager.list_tasks()
            workflow_tasks = workflow_data.get('tasks', [])
            self.logger.info(f"📊 工作流管理器任务数量: {len(workflow_tasks)}")
            if workflow_tasks:
                for task_info in workflow_tasks[:3]:  # 只显示前3个
                    task_id = task_info.get('task_id', 'unknown')
                    status = task_info.get('status', 'unknown')
                    self.logger.info(f"   - {task_id}: status={status}")

            # 合并任务列表
            all_tasks = {}

            # 创建workflow任务的字典，方便快速查找
            workflow_tasks_dict = {task['task_id']: task for task in workflow_tasks if 'task_id' in task}

            # 添加COM任务（优先级更高，先添加）
            for task_id, task_info in com_tasks.items():
                source = TaskSource.REMOTE_SCHEDULER.value if task_id.startswith('remote_') else TaskSource.HTTP_API.value

                # 如果是远程任务，需要提取原始任务ID并合并workflow信息
                if task_id.startswith('remote_'):
                    # 从内部任务ID提取原始任务ID: remote_{原始ID}_{随机} -> {原始ID}
                    parts = task_id.split('_')
                    if len(parts) >= 3:
                        # remote_764150237768974336_adc72e85 -> 764150237768974336
                        original_task_id = '_'.join(parts[1:-1])  # 支持task_id中包含下划线

                        # 查找对应的workflow任务信息
                        workflow_info = workflow_tasks_dict.get(original_task_id)

                        if workflow_info:
                            # 合并信息：使用原始ID，保留COM的状态和时间，添加workflow的详细信息
                            merged_info = {
                                **task_info,  # COM任务的基本信息（status, created_time等）
                                "task_id": original_task_id,  # 使用原始任务ID
                                "source": source,
                                # 添加workflow的详细信息
                                "progress": workflow_info.get('progress', 0),
                                "message": workflow_info.get('message', ''),
                                "url": workflow_info.get('url'),
                                "draft_name": workflow_info.get('draft_name'),
                            }
                            all_tasks[original_task_id] = merged_info
                            self.logger.debug(f"合并远程任务信息: {original_task_id} (COM: {task_id})")
                        else:
                            # 没有找到对应的workflow信息，使用原始ID但保持COM信息
                            all_tasks[original_task_id] = {
                                **task_info,
                                "task_id": original_task_id,
                                "source": source
                            }
                            self.logger.debug(f"远程任务无workflow信息: {original_task_id}")
                    else:
                        # 任务ID格式不正确，直接使用
                        all_tasks[task_id] = {
                            **task_info,
                            "source": source
                        }
                else:
                    # 非远程任务，直接添加
                    all_tasks[task_id] = {
                        **task_info,
                        "source": source
                    }

            # 添加工作流任务（去重：如果是远程任务对应的workflow任务，则跳过）
            for task_info in workflow_tasks:
                task_id = task_info.get('task_id')
                if task_id:
                    # 检查是否是远程任务对应的workflow任务（避免重复显示）
                    # 方法1: 检查是否在remote_task_mapping中
                    if task_id in self.remote_task_mapping:
                        self.logger.debug(f"跳过重复的workflow任务: {task_id} (在映射表中)")
                        continue

                    # 方法2: 检查COM任务中是否有包含此task_id的任务
                    is_duplicate = False
                    for com_task_id, com_task_info in com_tasks.items():
                        # 检查COM任务的kwargs中的task_id
                        kwargs = com_task_info.get('kwargs', {})
                        if kwargs.get('task_id') == task_id:
                            self.logger.debug(f"跳过重复的workflow任务: {task_id} (对应COM任务: {com_task_id})")
                            is_duplicate = True
                            break

                    if is_duplicate:
                        continue

                    # 不是远程任务，添加为HTTP_API来源
                    all_tasks[task_id] = {
                        **task_info,
                        "source": TaskSource.HTTP_API.value
                    }

            self.logger.info(f"📊 合并后总任务数: {len(all_tasks)}")
            if all_tasks:
                self.logger.info(f"📊 合并后任务ID列表: {list(all_tasks.keys())[:5]}")  # 只显示前5个

            return {
                "tasks": all_tasks,
                "summary": {
                    "total": len(all_tasks),
                    "pending": len([t for t in all_tasks.values() if t.get('status') == 'pending']),
                    "running": len([t for t in all_tasks.values() if t.get('status') == 'running']),
                    "completed": len([t for t in all_tasks.values() if t.get('status') == 'completed']),
                    "failed": len([t for t in all_tasks.values() if t.get('status') == 'failed'])
                },
                "timestamp": datetime.now().isoformat()
            }

        except Exception as e:
            self.logger.error(f"列出任务异常: {e}")
            import traceback
            self.logger.error(f"异常详情: {traceback.format_exc()}")
            return {
                "tasks": {},
                "summary": {"total": 0, "pending": 0, "running": 0, "completed": 0, "failed": 0},
                "timestamp": datetime.now().isoformat()
            }


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(level=logging.INFO)

    dispatcher = TaskDispatcher()

    # 测试HTTP任务
    http_task = {
        "url": "https://example.com/draft.json",
        "draft_name": "test_draft",
        "resolution": "1080p",
        "framerate": "30fps"
    }

    result = dispatcher.dispatch_http_task(http_task)
    print(f"HTTP任务分发结果: {result}")

    # 测试远程任务
    from scheduler_connector import RemoteTask

    remote_task = RemoteTask(
        id="remote_test_001",
        taskType="VIDEO_RENDER",
        taskParams={
            "draft_url": "https://example.com/remote_draft.json",
            "draft_name": "remote_test_draft"
        }
    )

    result = dispatcher.dispatch_remote_task(remote_task)
    print(f"远程任务分发结果: {result}")